package com.codeduck.kafka.producer;

import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.serialization.Serializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.Nullable;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

/**
 * @author: <a href="mailto:codeduck@163.com">Codeduck</a>
 * @date: 2023/6/23
 */
public class KafkaProducerFactory<K, V> implements ProducerFactory<K, V> {

    private static final Logger LOGGER = LoggerFactory.getLogger(KafkaProducerFactory.class);

    private final Map<String, Object> configs;

    private final AtomicInteger clientIdCounter = new AtomicInteger();

    private final ThreadLocal<Producer<K, V>> threadBoundProducers;

    private String clientIdPrefix;

    public KafkaProducerFactory(Map<String, Object> configs) {
        this(configs, () -> null, () -> null);
    }

    public KafkaProducerFactory(Map<String, Object> configs,
                                @Nullable Serializer<K> keySerializer,
                                @Nullable Serializer<V> valueSerializer) {
        this(configs, () -> keySerializer, () -> valueSerializer);
    }

    public KafkaProducerFactory(Map<String, Object> configs,
                                @Nullable Supplier<Serializer<K>> keySerializerSupplier,
                                @Nullable Supplier<Serializer<V>> valueSerializerSupplier) {
        this.configs = new HashMap<>(configs);
        this.threadBoundProducers = new ThreadLocal<>();
        if (this.clientIdPrefix == null && configs.get(ProducerConfig.CLIENT_ID_CONFIG) instanceof String) {
            this.clientIdPrefix = (String) configs.get(ProducerConfig.CLIENT_ID_CONFIG);
        }
    }

    @Override
    public Producer<K, V> createProducer() {
        if (threadBoundProducers.get() == null) {
            KafkaProducer<K, V> producer;
            if (clientIdPrefix == null) {
                producer = new KafkaProducer<>(this.configs);
            } else {
                Map<String, Object> newConfigs = new HashMap<>(this.configs);
                newConfigs.put(ProducerConfig.CLIENT_ID_CONFIG, this.clientIdPrefix + "_" + this.clientIdCounter.incrementAndGet());
                LOGGER.info("生产者Client_ID [{}", newConfigs.get(ProducerConfig.CLIENT_ID_CONFIG));
                producer = new KafkaProducer<>(newConfigs);
            }
            threadBoundProducers.set(producer);
        }
        return threadBoundProducers.get();
    }
}
